from comlibpy.xoption import XOption, XNone
from typing import Any, Callable, Optional

class Seq:

    def __init__(self, lst):
        """构建一个封装List(lst)的Seq封装类"""
        self.data = lst
        self.size = len(lst)
        self.st = 0
        self.ed = self.size
        self.step = 1
        self.is_mutable = False

    @staticmethod
    def fromElement(ele:Any, *sizes:int) -> 'Seq':
        """构建一个尺寸为sizes、元素都为ele的Seq"""
        def recur(ele, sizes):
            if len(sizes) == 1:
                return [ele for i in range(sizes[0])]
            else:
                return [recur(ele, sizes[1:]) for i in range(sizes[0])]
        return Seq(recur(ele, list(sizes)))

    @staticmethod
    def fromGenerator(gen:Callable[[list[int]], Any], *sizes:int) -> 'Seq':
        """构建一个尺寸为sizes的Seq。每个元素由构造函数gen产生"""
        def recur(gen:Callable[[list[int]], Any], position, sizes:list[int]):
            if len(sizes) == 1:
                return [gen(position+[i]) for i in range(sizes[0])]
            else:
                return [recur(gen, position+[i], sizes[1:]) for i in range(sizes[0])]
        return Seq(recur(gen, [], list(sizes)))


    def toList(self): return self.data

    def reaxis(self, dimMap:list[int]):
        """轴变换(维度变化)
        轴变换不改变数据结构，仅仅是内部存储格式的变换。比如：a[i][j]变为a[j][i]
        * dimMap: 轴变换映射表。dimMap[i]=j表示把当前tensor的第i个维度映射到新tensor的第j个维度
        * return: 返回新产生的tensor
        """

        #### 特殊情况处理
        if dimMap == [1,0]:
            nrow = self.size
            ncol = len(self.data[0])
            return Seq([[self.data[j][i] for j in range(nrow)] for i in range(ncol)])

        #### 计算变换涉及的原tensor各个维度的尺寸
        d_t = self.data 
        s_dim = len(dimMap)
        dimSize = [len(d_t)]
        for i in range(s_dim-1):
            d_t = d_t[0]
            dimSize.append( len(d_t) )

        #### 计算反向映射表
        revDimMap = [0 for i in range(s_dim)]
        for i, m in enumerate(dimMap):
            revDimMap[m] = i

        #### 计算映射后新tensor各个维度的尺寸
        newDimSize = [dimSize[revDimMap[i]] for i in range(s_dim)]

        def readData(pos:list[int]) -> Any:
            org_pos = [pos[m] for m in dimMap]
            return self.at(*org_pos)

        return Seq.fromGenerator(readData, *newDimSize)

    def at(self, *pos:int):
        """读取位置`pos`的数据。"""
        def recur( data, idx ):
            if len(idx) == 1:
                return data[idx[0]]
            else:
                return recur( data[idx[0]], idx[1:] )
        return recur(self.data, list(pos))

    def mutable(self, is_mutable:bool):
        """进入/退出 mutable模式"""
        self.is_mutable = is_mutable
        return self

    def sliceMode(self, st:Optional[int]=None, ed:Optional[int]=None, step:Optional[int]=None):
        """进入Slice处理模式。
        * st: 开始位置(include)。内部使用st%size(允许为负)。
        * ed: 结束位置(include)。内部使用ed%size(允许为负)
          - None: 原始数据长度 
          - other: 内部使用ed%size
        * step: 步长。
          - 注意：总是从st, st+step, st+2*step, ...。即使在reverse模式下
          - 大于0：递增Slice
          - 小于0：递减Slice
          - 当Step符号和st/ed不匹配时，函数会调换st/ed位置
        * 当st=ed=step=None时，表示恢复原始状态
        ## 常用例子：
          - sliceMode(): 恢复原始
          - sliceMode(step=-1): 在设定范围内反序
        """

        #### 每个参数有如下的可能：1) 设定；2）不动；3）恢复默认

        if st==None and ed==None and step==None:    #### 默认参数(特例)：恢复原状
            self.st = 0
            self.ed = self.size 
            self.step = 1

        elif st==None and ed==None:

            assert step!=0

            if self.step > 0 and step < 0:
                self.st, self.ed, self.step = self.ed - 1, self.st - 1, step
            elif self.step < 0 and step > 0:
                self.st, self.ed, self.step = self.ed + 1, self.st + 1, step
            else:
                self.step = step

        else:
            assert step!=0

            step = self.step if step==None else step

            if st == None:
                if self.step > 0: st = self.st 
                else: st = self.st + 1
            else:
                st = st % self.size 

            if ed == None:
                if self.step > 0: ed = self.ed 
                else: ed = self.ed + 1
            else:
                ed = ed % self.size
                if ed == 0: ed = self.size
            
            if (step < 0 and st < ed) or (step > 0 and st > ed):
                st, ed = ed, st

            if step > 0:
                self.st, self.ed, self.step = st, ed, step 
            else:
                self.st, self.ed, self.step = st - 1, ed - 1, step

        return self


    def shift(self, default=None):
        """从小到大元素依次移动位置"""

        if self.is_mutable:

            if self.size == 0: return self

            rpos = self.st 
            wpos = rpos + self.step 
            temp = self.data[rpos]
            while (wpos < self.ed and self.step > 0) or (self.step < 0 and wpos > self.ed):
                temp, self.data[wpos] = self.data[wpos], temp
                rpos = wpos 
                wpos = wpos + self.step

            self.data[self.st] = default if default!=None else temp
            return self

        else:

            if self.size == 0: return Seq([])

            res = [self.data[0]]
            rpos = self.st 
            wpos = rpos + self.step 
            while (wpos < self.ed and self.step > 0) or (self.step < 0 and wpos > self.ed):
                res.append(self.data[rpos])
                rpos = wpos 
                wpos = wpos + self.step

            res[0] = default if default!=None else self.data[wpos-self.step]

            return Seq(res)


    def find2(self, criteria:Callable[[Any,int], bool]) -> XOption:
        """查找满足条件`criteria`的所有元素"""

        for i in range(self.st, self.ed, self.step):
            res = criteria(self.data[i], i)
            if not res.isEmpty(): return res

        return XNone


    def find(self, criteria, frm=0, reverse=False):

        sop = frm if frm>=0 else self.size+frm
        eop = 0 if reverse else self.size
        step = -1 if reverse else 1

        for i in range(sop, eop, step):
            res = criteria(self.data[i], i)
            if not res.isEmpty(): return res

        return XNone


    def count(self, criteria:Callable[[Any,int], bool]) -> int:
        """统计满足条件`criteria`的元素的个数"""
        res = 0

        for i in range(self.st, self.ed, self.step):
            e = self.data[i]
            if criteria(e,i):
                res = res + 1
        return res

    def index2(self, criteria:Callable[[Any,int], bool], num:int=1) -> list[int]:
        """查找满足条件`criteria`的元素的位置。
        * criteria: 条件
          - e: 元素
          - idx: 位置
          - return: 是否满足条件
        * num: 查找最大个数
          - >0: 查找个数
          - <0: 查找所有满足条件的元素
        * return: 返回满足条件的所有元素的位置信息
        """
        res = []

        for i in range(self.st, self.ed, self.step):
            e = self.data[i]
            if criteria(e,i):
                res.append(i)
                if len(res) >= num:
                    return res
        return res

    def index(self, criteria, num=1, dir=0):
        res = []
        if num<0: num=self.size
        for i,e in enumerate(self.data):
            if criteria(e,i):
                res.append(i)
                if len(res) >= num:
                    return res
        return res
